import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QLabel, QSlider, QGroupBox, QFrame,
                             QPushButton, QLineEdit)
from PySide6.QtCore import Qt
import cv2

# 设置matplotlib使用Qt6后端
plt.switch_backend('QtAgg')

# 靶标坐标系下的6个点坐标（单位：毫米）
target_points = np.array([
    [0,   0,   0],    # P1
    [500, 0,   0],    # P2
    [500, 500, 0],    # P3
    [0,   500, 0],    # P4
    [-250,   250,   -300],  # P5
    [750, 250, -300],  # P6
], dtype=np.float64)  # 确保使用float64类型

def read_intrinsic_from_txt(filename):
    """从txt文件中读取相机内参数据"""
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            
        # 初始化参数
        fx = fy = cx = cy = None
        fx_true = fy_true = cx_true = cy_true = None
        
        # 遍历文件内容查找参数
        for line in lines:
            line = line.strip()
            if line.startswith('fx:'):
                fx = float(line.split(':')[1].strip())
            elif line.startswith('fy:'):
                fy = float(line.split(':')[1].strip())
            elif line.startswith('cx:'):
                cx = float(line.split(':')[1].strip())
            elif line.startswith('cy:'):
                cy = float(line.split(':')[1].strip())
            elif line.startswith('fx_true:'):
                fx_true = float(line.split(':')[1].strip())
            elif line.startswith('fy_true:'):
                fy_true = float(line.split(':')[1].strip())
            elif line.startswith('cx_true:'):
                cx_true = float(line.split(':')[1].strip())
            elif line.startswith('cy_true:'):
                cy_true = float(line.split(':')[1].strip())
        
        # 检查是否所有参数都已找到
        if all(param is not None for param in [fx, fy, cx, cy]):
            return np.array([fx, fy, cx, cy]), np.array([fx_true, fy_true, cx_true, cy_true])
        else:
            print("未找到完整的相机内参数据")
            return None, None
            
    except FileNotFoundError:
        print("未找到相机参数文件")
        return None, None
    except Exception as e:
        print(f"读取相机参数出错: {str(e)}")
        return None, None

def project_points(points_3d, R, t, camera_matrix):
    """投影3D点到2D图像平面"""
    # 确保输入数据格式正确
    points_3d = points_3d.astype(np.float32)
    R = R.astype(np.float32)
    t = t.astype(np.float32).reshape(3, 1)  # 确保t是3x1的列向量
    camera_matrix = camera_matrix.astype(np.float32)
    
    # 将点转换到相机坐标系
    points_cam = R @ points_3d.T + t
    
    # 投影到图像平面
    points_2d = camera_matrix @ points_cam
    points_2d = points_2d[:2] / points_2d[2]
    
    return points_2d.T

def draw_target_points(ax, points_2d, distance):
    """在matplotlib图像上绘制靶标点和点号"""
    # 清除之前的图像
    ax.clear()
    
    # 设置图像范围
    ax.set_xlim(0, 1920)
    ax.set_ylim(1080, 0)  # 反转Y轴以匹配图像坐标系
    ax.set_aspect('equal')
    
    # 绘制点和点号
    for i, point in enumerate(points_2d):
        # 绘制点
        ax.plot(point[0], point[1], 'wo', markersize=8)  # 白色点
        # 添加点号标签
        ax.text(point[0] + 20, point[1] + 20, f'P{i+1}', 
                color='white', fontsize=12, fontweight='bold')
    
    # 添加距离信息
    ax.text(10, 30, f'Distance: {distance/1000:.2f}m', 
            color='white', fontsize=12, fontweight='bold')
    
    # 设置背景为黑色
    ax.set_facecolor('black')
    
    # 隐藏坐标轴
    ax.axis('off')
    
    # 刷新图像
    ax.figure.canvas.draw()

def solve_pnp(points_3d, points_2d, camera_matrix):
    """使用PnP算法求解靶标位置和姿态"""
    # 确保输入数据格式正确
    points_3d = points_3d.astype(np.float32)
    points_2d = points_2d.astype(np.float32)
    camera_matrix = camera_matrix.astype(np.float32)
    
    # 使用EPnP算法求解
    success, rvec, tvec = cv2.solvePnP(points_3d, points_2d, camera_matrix, None, flags=cv2.SOLVEPNP_EPNP)
    
    if success:
        # 将旋转向量转换为旋转矩阵
        R, _ = cv2.Rodrigues(rvec)
        return R, tvec
    else:
        return None, None

class TargetControlWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("靶标控制界面")
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建主窗口部件
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        
        # 创建主布局
        layout = QHBoxLayout(main_widget)
        
        # 左侧：相机内参显示
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        
        # 读取相机内参
        intrinsic_params, true_params = read_intrinsic_from_txt("based on pts/camera_intrinsics.txt")
        if intrinsic_params is not None and true_params is not None:
            fx, fy, cx, cy = intrinsic_params
            fx_true, fy_true, cx_true, cy_true = true_params
            
            # 使用标定值创建相机矩阵
            self.camera_matrix = np.array([
                [fx, 0, cx],
                [0, fy, cy],
                [0, 0, 1]
            ])
            
            # 使用真值创建相机矩阵（仅用于显示）
            self.true_camera_matrix = np.array([
                [fx_true, 0, cx_true],
                [0, fy_true, cy_true],
                [0, 0, 1]
            ])
            
            # 计算焦距（单位：毫米）
            pixel_size = 0.0029  # 像素尺寸（毫米）
            focal_length_x = fx * pixel_size
            focal_length_y = fy * pixel_size
            focal_length_x_true = fx_true * pixel_size
            focal_length_y_true = fy_true * pixel_size
            
            # 显示相机内参
            intrinsic_group = QGroupBox("相机内参")
            intrinsic_layout = QVBoxLayout()
            
            # 创建水平布局用于左右排布
            params_layout = QHBoxLayout()
            
            # 创建左侧真值布局
            true_layout = QVBoxLayout()
            true_label = QLabel("真值：")
            true_label.setStyleSheet("color: #FFA500; font-weight: bold;")  # 橙色
            true_layout.addWidget(true_label)
            
            # 添加真值信息（橙色）
            true_style = "color: #FFA500;"
            true_layout.addWidget(QLabel(f"焦距 X: {focal_length_x_true:.2f} mm", styleSheet=true_style))
            true_layout.addWidget(QLabel(f"焦距 Y: {focal_length_y_true:.2f} mm", styleSheet=true_style))
            true_layout.addWidget(QLabel(f"fx = {fx_true:.2f} 像素", styleSheet=true_style))
            true_layout.addWidget(QLabel(f"fy = {fy_true:.2f} 像素", styleSheet=true_style))
            true_layout.addWidget(QLabel(f"cx = {cx_true:.2f} 像素", styleSheet=true_style))
            true_layout.addWidget(QLabel(f"cy = {cy_true:.2f} 像素", styleSheet=true_style))
            
            # 创建中间标定值布局
            cal_layout = QVBoxLayout()
            cal_label = QLabel("标定值：")
            cal_label.setStyleSheet("color: #0000FF; font-weight: bold;")  # 蓝色
            cal_layout.addWidget(cal_label)
            
            # 添加标定值信息（蓝色）
            cal_style = "color: #0000FF;"
            cal_layout.addWidget(QLabel(f"焦距 X: {focal_length_x:.2f} mm", styleSheet=cal_style))
            cal_layout.addWidget(QLabel(f"焦距 Y: {focal_length_y:.2f} mm", styleSheet=cal_style))
            cal_layout.addWidget(QLabel(f"fx = {fx:.2f} 像素", styleSheet=cal_style))
            cal_layout.addWidget(QLabel(f"fy = {fy:.2f} 像素", styleSheet=cal_style))
            cal_layout.addWidget(QLabel(f"cx = {cx:.2f} 像素", styleSheet=cal_style))
            cal_layout.addWidget(QLabel(f"cy = {cy:.2f} 像素", styleSheet=cal_style))
            
            # 创建右侧误差布局
            error_layout = QVBoxLayout()
            error_label = QLabel("误差：")
            error_label.setStyleSheet("color: #FF0000; font-weight: bold;")  # 红色
            error_layout.addWidget(error_label)
            
            # 计算并添加误差信息（红色）
            error_style = "color: #FF0000;"
            error_layout.addWidget(QLabel(f"Δ焦距 X: {focal_length_x - focal_length_x_true:.2f} mm", styleSheet=error_style))
            error_layout.addWidget(QLabel(f"Δ焦距 Y: {focal_length_y - focal_length_y_true:.2f} mm", styleSheet=error_style))
            error_layout.addWidget(QLabel(f"Δfx = {fx - fx_true:.2f} 像素", styleSheet=error_style))
            error_layout.addWidget(QLabel(f"Δfy = {fy - fy_true:.2f} 像素", styleSheet=error_style))
            error_layout.addWidget(QLabel(f"Δcx = {cx - cx_true:.2f} 像素", styleSheet=error_style))
            error_layout.addWidget(QLabel(f"Δcy = {cy - cy_true:.2f} 像素", styleSheet=error_style))
            
            # 将三个布局添加到水平布局中
            params_layout.addLayout(true_layout)
            params_layout.addLayout(cal_layout)
            params_layout.addLayout(error_layout)
            
            # 将水平布局添加到主布局中
            intrinsic_layout.addLayout(params_layout)
            
            intrinsic_group.setLayout(intrinsic_layout)
            left_layout.addWidget(intrinsic_group)
        
        # 添加靶标位置和姿态显示组
        pose_group = QGroupBox("靶标位置和姿态")
        pose_layout = QVBoxLayout()
        
        # 创建水平布局用于左右排布
        pose_params_layout = QHBoxLayout()
        
        # 创建左侧真值布局
        pose_true_layout = QVBoxLayout()
        pose_true_label = QLabel("真值：")
        pose_true_label.setStyleSheet("color: #FFA500; font-weight: bold;")  # 橙色
        pose_true_layout.addWidget(pose_true_label)
        
        # 创建位置和姿态标签
        self.position_labels = {
            'x': QLabel("X: 0.00 mm"),
            'y': QLabel("Y: 0.00 mm"),
            'z': QLabel("Z: 0.00 mm")
        }
        self.rotation_labels = {
            'roll': QLabel("Roll: 0.00°"),
            'pitch': QLabel("Pitch: 0.00°"),
            'yaw': QLabel("Yaw: 0.00°")
        }
        
        # 添加位置标签
        for label in self.position_labels.values():
            pose_true_layout.addWidget(label)
        
        # 添加姿态标签
        for label in self.rotation_labels.values():
            pose_true_layout.addWidget(label)
        
        # 创建中间计算值布局
        pose_cal_layout = QVBoxLayout()
        pose_cal_label = QLabel("计算值：")
        pose_cal_label.setStyleSheet("color: #0000FF; font-weight: bold;")  # 蓝色
        pose_cal_layout.addWidget(pose_cal_label)
        
        # 创建计算值标签
        self.calc_position_labels = {
            'x': QLabel("X: 0.00 mm"),
            'y': QLabel("Y: 0.00 mm"),
            'z': QLabel("Z: 0.00 mm")
        }
        self.calc_rotation_labels = {
            'roll': QLabel("Roll: 0.00°"),
            'pitch': QLabel("Pitch: 0.00°"),
            'yaw': QLabel("Yaw: 0.00°")
        }
        
        # 添加计算值标签
        for label in self.calc_position_labels.values():
            pose_cal_layout.addWidget(label)
        
        for label in self.calc_rotation_labels.values():
            pose_cal_layout.addWidget(label)
            
        # 创建右侧误差布局
        pose_error_layout = QVBoxLayout()
        pose_error_label = QLabel("误差：")
        pose_error_label.setStyleSheet("color: #FF0000; font-weight: bold;")  # 红色
        pose_error_layout.addWidget(pose_error_label)
        
        # 创建误差标签
        self.error_position_labels = {
            'x': QLabel("ΔX: 0.00 mm"),
            'y': QLabel("ΔY: 0.00 mm"),
            'z': QLabel("ΔZ: 0.00 mm")
        }
        self.error_rotation_labels = {
            'roll': QLabel("ΔRoll: 0.00°"),
            'pitch': QLabel("ΔPitch: 0.00°"),
            'yaw': QLabel("ΔYaw: 0.00°")
        }
        
        # 为所有误差标签设置红色
        error_style = "color: #FF0000;"
        for label in self.error_position_labels.values():
            label.setStyleSheet(error_style)
        for label in self.error_rotation_labels.values():
            label.setStyleSheet(error_style)
        
        # 添加误差标签
        for label in self.error_position_labels.values():
            pose_error_layout.addWidget(label)
        
        for label in self.error_rotation_labels.values():
            pose_error_layout.addWidget(label)
        
        # 将三个布局添加到水平布局中
        pose_params_layout.addLayout(pose_true_layout)
        pose_params_layout.addLayout(pose_cal_layout)
        pose_params_layout.addLayout(pose_error_layout)
        
        # 将水平布局添加到主布局中
        pose_layout.addLayout(pose_params_layout)
        
        # 创建解算按钮和输入框的水平布局
        solve_layout = QHBoxLayout()
        
        # 添加解算按钮
        solve_button = QPushButton("解算")
        solve_button.clicked.connect(self.solve_pose)
        solve_layout.addWidget(solve_button)
        
        # 添加质心标准差输入框
        centroid_std_layout = QHBoxLayout()
        centroid_std_label = QLabel("质心标准差(像素):")
        self.centroid_std_input = QLineEdit("0.5")
        self.centroid_std_input.setFixedWidth(60)
        centroid_std_layout.addWidget(centroid_std_label)
        centroid_std_layout.addWidget(self.centroid_std_input)
        solve_layout.addLayout(centroid_std_layout)
        
        # 添加靶标标定标准差输入框
        target_std_layout = QHBoxLayout()
        target_std_label = QLabel("靶标标定标准差(mm):")
        self.target_std_input = QLineEdit("1.0")
        self.target_std_input.setFixedWidth(60)
        target_std_layout.addWidget(target_std_label)
        target_std_layout.addWidget(self.target_std_input)
        solve_layout.addLayout(target_std_layout)
        
        # 将解算布局添加到主布局
        pose_layout.addLayout(solve_layout)
        
        pose_group.setLayout(pose_layout)
        left_layout.addWidget(pose_group)
        
        # 右侧：控制滑块
        control_group = QGroupBox("靶标控制")
        control_layout = QVBoxLayout()
        
        # 创建滑块
        self.sliders = {}
        slider_params = {
            'x': (-10000, 10000, 0, "X位置 (mm)"),
            'y': (-10000, 10000, 0, "Y位置 (mm)"),
            'z': (10000, 200000, 50000, "Z位置 (mm)"),
            'roll': (-45, 45, 0, "Roll (度)"),
            'pitch': (-45, 45, 0, "Pitch (度)"),
            'yaw': (-45, 45, 0, "Yaw (度)")
        }
        
        for name, (min_val, max_val, default, label) in slider_params.items():
            slider_layout = QHBoxLayout()
            slider_layout.addWidget(QLabel(label))
            slider = QSlider(Qt.Horizontal)
            slider.setMinimum(min_val)
            slider.setMaximum(max_val)
            slider.setValue(default)
            slider.valueChanged.connect(self.update_projection)
            self.sliders[name] = slider
            slider_layout.addWidget(slider)
            control_layout.addLayout(slider_layout)
        
        control_group.setLayout(control_layout)
        left_layout.addWidget(control_group)
        
        # 添加左侧面板到主布局
        layout.addWidget(left_panel)
        
        # 右侧：matplotlib图像显示
        self.figure = plt.figure(figsize=(8, 6))
        self.canvas = FigureCanvas(self.figure)
        self.ax = self.figure.add_subplot(111)
        self.ax.set_facecolor('black')
        self.figure.patch.set_facecolor('black')
        layout.addWidget(self.canvas)
        
        # 初始化投影
        self.update_projection()
    
    def solve_pose(self):
        """解算靶标位置和姿态"""
        # 获取当前投影点
        x = float(self.sliders['x'].value())
        y = float(self.sliders['y'].value())
        z = float(self.sliders['z'].value())
        roll = np.radians(float(self.sliders['roll'].value()))
        pitch = np.radians(float(self.sliders['pitch'].value()))
        yaw = np.radians(float(self.sliders['yaw'].value()))
        
        # 获取标准差输入值
        try:
            centroid_std = float(self.centroid_std_input.text())
            target_std = float(self.target_std_input.text())
        except ValueError:
            # 如果输入无效，使用默认值
            centroid_std = 0.5
            target_std = 1.0
            self.centroid_std_input.setText("0.5")
            self.target_std_input.setText("1.0")
        
        # 计算旋转矩阵
        Rx = np.array([
            [1, 0, 0],
            [0, np.cos(roll), -np.sin(roll)],
            [0, np.sin(roll), np.cos(roll)]
        ], dtype=np.float64)
        
        Ry = np.array([
            [np.cos(pitch), 0, np.sin(pitch)],
            [0, 1, 0],
            [-np.sin(pitch), 0, np.cos(pitch)]
        ], dtype=np.float64)
        
        Rz = np.array([
            [np.cos(yaw), -np.sin(yaw), 0],
            [np.sin(yaw), np.cos(yaw), 0],
            [0, 0, 1]
        ], dtype=np.float64)
        
        R = Rz @ Ry @ Rx
        t = np.array([x, y, z], dtype=np.float64)
        
        # 使用真值相机矩阵进行投影
        projected_points = project_points(target_points, R, t, self.true_camera_matrix)
        
        # 添加质心标准差到投影点坐标
        noise = np.random.normal(0, centroid_std, projected_points.shape)
        projected_points_noisy = projected_points + noise
        
        # 添加靶标标定标准差到靶标点坐标
        target_points_noisy = target_points.copy()
        # 将靶标标定标准差从毫米转换为像素（使用平均焦距）
        avg_focal_length = (self.camera_matrix[0,0] + self.camera_matrix[1,1]) / 2
        target_std_pixels = target_std * avg_focal_length / z
        noise_3d = np.random.normal(0, target_std_pixels, target_points.shape)
        target_points_noisy = target_points_noisy + noise_3d  # 使用显式赋值而不是+=操作符
        
        # 使用标定值相机矩阵进行解算
        R_calc, t_calc = solve_pnp(target_points_noisy, projected_points_noisy, self.camera_matrix)
        
        if R_calc is not None and t_calc is not None:
            # 从旋转矩阵中提取欧拉角
            roll_calc = np.arctan2(R_calc[2,1], R_calc[2,2])
            pitch_calc = np.arctan2(-R_calc[2,0], np.sqrt(R_calc[2,1]**2 + R_calc[2,2]**2))
            yaw_calc = np.arctan2(R_calc[1,0], R_calc[0,0])
            
            # 将NumPy数组转换为Python标量值
            t_calc = t_calc.flatten()  # 确保t_calc是一维数组
            
            # 更新计算值显示
            self.calc_position_labels['x'].setText(f"X: {float(t_calc[0]):.2f} mm")
            self.calc_position_labels['y'].setText(f"Y: {float(t_calc[1]):.2f} mm")
            self.calc_position_labels['z'].setText(f"Z: {float(t_calc[2]):.2f} mm")
            self.calc_rotation_labels['roll'].setText(f"Roll: {float(np.degrees(roll_calc)):.2f}°")
            self.calc_rotation_labels['pitch'].setText(f"Pitch: {float(np.degrees(pitch_calc)):.2f}°")
            self.calc_rotation_labels['yaw'].setText(f"Yaw: {float(np.degrees(yaw_calc)):.2f}°")
            
            # 计算误差
            pos_error = t_calc - np.array([x, y, z])
            rot_error = np.array([
                np.degrees(roll_calc) - np.degrees(roll),
                np.degrees(pitch_calc) - np.degrees(pitch),
                np.degrees(yaw_calc) - np.degrees(yaw)
            ])
            
            # 更新误差显示
            self.error_position_labels['x'].setText(f"ΔX: {float(pos_error[0]):.2f} mm")
            self.error_position_labels['y'].setText(f"ΔY: {float(pos_error[1]):.2f} mm")
            self.error_position_labels['z'].setText(f"ΔZ: {float(pos_error[2]):.2f} mm")
            self.error_rotation_labels['roll'].setText(f"ΔRoll: {float(rot_error[0]):.2f}°")
            self.error_rotation_labels['pitch'].setText(f"ΔPitch: {float(rot_error[1]):.2f}°")
            self.error_rotation_labels['yaw'].setText(f"ΔYaw: {float(rot_error[2]):.2f}°")
    
    def update_projection(self):
        # 获取滑块值
        x = self.sliders['x'].value()
        y = self.sliders['y'].value()
        z = self.sliders['z'].value()
        roll = np.radians(self.sliders['roll'].value())
        pitch = np.radians(self.sliders['pitch'].value())
        yaw = np.radians(self.sliders['yaw'].value())
        
        # 更新位置和姿态显示
        self.position_labels['x'].setText(f"X: {x:.2f} mm")
        self.position_labels['y'].setText(f"Y: {y:.2f} mm")
        self.position_labels['z'].setText(f"Z: {z:.2f} mm")
        self.rotation_labels['roll'].setText(f"Roll: {np.degrees(roll):.2f}°")
        self.rotation_labels['pitch'].setText(f"Pitch: {np.degrees(pitch):.2f}°")
        self.rotation_labels['yaw'].setText(f"Yaw: {np.degrees(yaw):.2f}°")
        
        # 计算旋转矩阵
        Rx = np.array([
            [1, 0, 0],
            [0, np.cos(roll), -np.sin(roll)],
            [0, np.sin(roll), np.cos(roll)]
        ])
        
        Ry = np.array([
            [np.cos(pitch), 0, np.sin(pitch)],
            [0, 1, 0],
            [-np.sin(pitch), 0, np.cos(pitch)]
        ])
        
        Rz = np.array([
            [np.cos(yaw), -np.sin(yaw), 0],
            [np.sin(yaw), np.cos(yaw), 0],
            [0, 0, 1]
        ])
        
        R = Rz @ Ry @ Rx
        t = np.array([x, y, z])
        
        # 使用真值相机矩阵进行投影显示
        projected_points = project_points(target_points, R, t, self.true_camera_matrix)
        
        # 绘制投影
        draw_target_points(self.ax, projected_points, z)

def main():
    app = QApplication([])
    window = TargetControlWindow()
    window.show()
    app.exec()

if __name__ == "__main__":
    main()



